package chash

import (
	"encoding/binary"
	"errors"
	blake2b "github.com/minio/blake2b-simd"
	"sort"
	"strconv"
)

// The Add and Get interface not goroutine safe
const (
	replicationFactor = 64
)

type CHash struct {
	hosts     map[uint64]string
	hostsMap  map[string]bool
	sortedSet []uint64
}

func (c *CHash) hash(key string) uint64 {
	out := blake2b.Sum512([]byte(key))
	return binary.LittleEndian.Uint64(out[:])
}

func (c *CHash) search(key uint64) int {
	idx := sort.Search(len(c.sortedSet), func(i int) bool {
		return c.sortedSet[i] >= key
	})

	if idx >= len(c.sortedSet) {
		idx = 0
	}

	return idx
}

func NewCHash() *CHash {
	return &CHash{
		hosts:     make(map[uint64]string),
		hostsMap:  make(map[string]bool),
		sortedSet: make([]uint64, 0, 1024),
	}
}

func (c *CHash) Add(host string, weight int) {
	if c.hostsMap[host] {
		return
	}

	c.hostsMap[host] = true

	replace := replicationFactor * weight
	if replace < 1 {
		return
	}

	for i := 0; i < replace; i++ {
		h := c.hash(host + "#" + strconv.Itoa(i))
		c.hosts[h] = host
		c.sortedSet = append(c.sortedSet, h)
	}

	sort.Slice(c.sortedSet, func(i int, j int) bool {
		if c.sortedSet[i] < c.sortedSet[j] {
			return true
		}
		return false
	})
}

var (
	ErrNoHosts = errors.New("no hosts")
)

func (c *CHash) Get(key string) (string, error) {
	if len(c.hosts) == 0 {
		return "", ErrNoHosts
	}

	h := c.hash(key)
	idx := c.search(h)
	return c.hosts[c.sortedSet[idx]], nil
}
